import lstm
import time
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from keras.models import load_model

TIME_SERIES_OUTPUT_DATA = '/home/pan/forewarning_system/time_series/output/time_series.pkl'


def read_pd_pkl(path_data):
    return pd.read_pickle(path_data)


def plot_results(predicted_data, true_data):
    fig = plt.figure(facecolor='white', figsize=(80, 40))
    ax = fig.add_subplot(111)
    ax.plot(true_data, 'ro', label='True Data', alpha=0.5)
    plt.plot(predicted_data, 'bo', label='Prediction', alpha=0.5)
    plt.legend()
    plt.savefig('/home/pan/disk_windows_share/潘永灿/定子温度_v2/time/linear.jpg', dpi=100)

    # plt.show()
    print('plot_results finished ...')


def plot_scatter(predicted_data, true_data):
    plt.figure(figsize=(80, 40))
    plt.scatter(true_data, predicted_data)
    plt.plot([0, 10, 20, 30, 40, 50], [0, 10, 20, 30, 40, 50])
    plt.grid('on')
    plt.savefig('/home/pan/disk_windows_share/潘永灿/定子温度_v2/time/scatter.jpg', dpi=100)
    print('plot_scatter finished ...')


def plot_results_multiple(predicted_data, true_data, prediction_len):
    fig = plt.figure(facecolor='white')
    ax = fig.add_subplot(111)
    ax.plot(true_data, label='True Data')
    # Pad the list of predictions to shift it in the graph to it's correct start
    for i, data in enumerate(predicted_data):
        padding = [None for p in range(i * prediction_len)]
        plt.plot(padding + data, label='Prediction')
        plt.legend()
    plt.show()


# Main Run Thread
if __name__ == '__main__':
    global_start_time = time.time()
    epochs = 1
    data = read_pd_pkl(TIME_SERIES_OUTPUT_DATA)
    X_train = data.ix[:, 0:6]
    y_train = data.ix[:, 6]
    X_train = np.array(X_train)
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
    X_test = X_train
    y_test = y_train
    # train
    print('> Data Loaded. Compiling...')

    # model = lstm.build_model([1, 50, 100, 1])
    model = lstm.build_model([1, 6, 200, 1])

    model.fit(
        X_train,
        y_train,
        batch_size=512,
        nb_epoch=epochs,
        validation_split=0.05)

    # predictions = lstm.predict_sequences_multiple(model, X_test, seq_len, 50)
    # predictions = lstm.predict_sequences_multiple(model, X_test, seq_len, 6)
    # predicted = lstm.predict_sequence_full(model, X_test, seq_len)
    print('save model ...')
    model.save('model')

    # plot graph
    print('load model ...')
    model = load_model('model')
    print('predict ...')
    predicted = lstm.predict_point_by_point(model, X_test)

    print('Training duration (s) : ', time.time() - global_start_time)
    # plot_results_multiple(predictions, y_test, 50)
    plot_results(predicted, y_test)
    plot_scatter(predicted, y_test)
